import requests
import base64
import sys
import os
from PIL import Image
from io import BytesIO
from main import IMG_API_KEY,IMG_API_URL,IMG_SAVE_PATH,IMG_MODEL_NAME


class img_model:
    def __init__(self, api_key, end_point, img_model_name, img_save_path):
        self.IMG_API_KEY = api_key
        self.END_POINT = end_point
        self.IMG_MODEL_NAME = img_model_name
        self.PATH = img_save_path

    def save_image(self, url, save_path):
        try:
            response = requests.get(url, stream=True)
            if response.status_code == 200:
                # 自动创建目录
                os.makedirs(os.path.dirname(save_path), exist_ok=True)

                # 写入文件
                with open(save_path, 'wb') as f:
                    for chunk in response.iter_content(1024):
                        f.write(chunk)
                print(f"图片已保存至：{save_path}")
            else:
                print(f"下载失败，状态码：{response.status_code}")
        except Exception as e:
            print(f"发生错误：{str(e)}")

    def img_gen(self, img_size, prompt, img):
        headers = {
            "Authorization": f"Bearer {self.IMG_API_KEY}",
            "Content-Type": "application/json"
        }
        if bool(img):
            payload = {
                "model_name": self.IMG_MODEL_NAME,
                "prompt": prompt,
                "aspect_ratio": img_size,
                "n": 4
            }
        else:
            payload = {
                "model_name": self.IMG_MODEL_NAME,
                "prompt": prompt,
                "image": base64.b64encode(img).decode("utf-8"),
                "aspect_ratio": img_size,
                "n": 4
            }
        response = requests.post(self.END_POINT, headers=headers, json=payload)
        if response.status_code == 200:
            image_url = response.json()['data']['task_result']['images']['url']
            self.save_image(image_url, self.PATH)


img_client = img_model(IMG_API_KEY, IMG_API_URL, IMG_MODEL_NAME, IMG_SAVE_PATH)

if len(sys.argv) == 4 or len(sys.argv) == 3:
    if len(sys.argv) == 3:
        Prompt = sys.argv[1]
        Img_size = sys.argv[2]
        buffer = None
    else:
        Prompt = sys.argv[1]
        Img_size = sys.argv[2]
        buffer = BytesIO()
        Img = Image.open(sys.argv[3])
        Img.save(buffer, format="JPEG")
        buffer.seek(0)
    img_client.img_gen(img_size=Img_size, prompt=Prompt, img=buffer)
else:
    print("false")
